"""
wrapper for face detection, align and recognition
"""
import cv2
import numpy as np
from .scrfd import SCRFD
from .alignment import FaceAligner
from .arcface import Arcface


def get_detection_model(model_name, model_weights):
    if model_name == 'scrfd':
        model = SCRFD(model_file=model_weights)
        model.prepare(ctx_id=0)
        return model
    else:
        raise NotImplementedError


def get_recognition_model(model_name, model_weights):
    if model_name == 'arcface':
        return Arcface(model_weights)
    else:
        raise NotImplementedError


class FacePipeline(object):
    def __init__(self, config):
        # load models
        self.detector = get_detection_model(config.DET_MODEL, config.DET_WEIGHTS)
        self.face_aligner = FaceAligner(112, 'arcface')
        self.feature_extractor = get_recognition_model(config.REC_MODEL, config.REC_WEIGHTS) 

    def normalize_image(self, image):
        if image is None: return None
        if len(image.shape) == 2:
            image = image[:, :, np.newaxis]
        if image.shape[-1] == 1:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        elif image.shape[-1] == 4:
            image = image[:, :, :3]
        return image

    def run(self, image_path):
        # image path may contain chinese characters.
        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
        image = self.normalize_image(image)
        if image is None:
            return None
        bboxes, kpss = self.detector.detect(image, 0.5, input_size=(640, 640))
        ndets = len(bboxes)
        
        features = []
        for i in range(ndets):
            image_crop = self.face_aligner.norm_crop(image, kpss[i])
            feat = self.feature_extractor.predict_image(image_crop)
            # normalize face feature
            length = np.linalg.norm(feat)
            feat = feat / length
            features.append(feat)
        return features
    
    def get_face_crops(self, image_path):
        """
        generate aligned face crops for feature extraction
        """
        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
        image = self.normalize_image(image)
        if image is None:
            return []
        
        crops = []
        bboxes, kpss = self.detector.detect(image, 0.5, input_size=(640, 640))
        ndets = len(bboxes)
        for i in range(ndets):
            image_crop = self.face_aligner.norm_crop(image, kpss[i])
            crops.append(image_crop)
        
        return crops
    
    def get_face_feature(self, image_path):
        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
        feat = self.feature_extractor.predict_image(image)
        return feat


def get_model(model_cfg):
    return FacePipeline(model_cfg)


if __name__ == '__main__':
    model = FacePipeline()
    image_path = 'data/test/test.jpg'
    feats = model.run(image_path)
    print(feats)
